# src/security.py
import subprocess
import json
import os
import shutil
from typing import Dict, List, Optional
from datetime import datetime


class SecurityScanner:
    """容器镜像安全扫描器"""
    
    def __init__(self):
        self.trivy_available = self._check_trivy_available()
    
    def _check_trivy_available(self) -> bool:
        """检查 trivy 是否可用"""
        try:
            result = subprocess.run(
                ["trivy", "--version"], 
                capture_output=True, 
                text=True, 
                timeout=10
            )
            return result.returncode == 0
        except (subprocess.TimeoutExpired, FileNotFoundError):
            return False
    
    def _install_trivy_suggestion(self) -> str:
        """返回安装 trivy 的建议"""
        return """
要使用安全扫描功能，请先安装 trivy:

# Linux (curl)
curl -sfL https://raw.githubusercontent.com/aquasecurity/trivy/main/contrib/install.sh | sh -s -- -b /usr/local/bin

# macOS (Homebrew)
brew install trivy

# 或者使用 Docker 运行 (无需本地安装)
docker run --rm -v /var/run/docker.sock:/var/run/docker.sock aquasec/trivy:latest image <镜像名>
"""

    def scan_image(self, image_name: str, output_format: str = "json") -> Optional[Dict]:
        """
        扫描镜像安全漏洞
        
        Args:
            image_name: 镜像名称
            output_format: 输出格式 (json/table)
            
        Returns:
            扫描结果字典，失败时返回 None
        """
        if not self.trivy_available:
            print("⚠️  Trivy 未安装，跳过安全扫描")
            print(self._install_trivy_suggestion())
            return None
        
        try:
            cmd = [
                "trivy", "image",
                "--format", output_format,
                "--severity", "HIGH,CRITICAL,MEDIUM,LOW",
                "--no-progress",
                image_name
            ]
            
            print(f"   执行扫描命令: trivy image {image_name}")
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=300  # 5分钟超时
            )
            
            if result.returncode != 0:
                print(f"   扫描失败: {result.stderr}")
                return None
            
            if output_format == "json":
                return json.loads(result.stdout) if result.stdout else {}
            else:
                return {"raw_output": result.stdout}
                
        except subprocess.TimeoutExpired:
            print("   扫描超时（5分钟）")
            return None
        except json.JSONDecodeError as e:
            print(f"   JSON 解析失败: {e}")
            return None
        except Exception as e:
            print(f"   扫描过程中发生错误: {e}")
            return None

    def generate_security_report(self, scan_result: Dict, image_name: str, output_path: str) -> bool:
        """
        生成安全报告
        
        Args:
            scan_result: 扫描结果
            image_name: 镜像名称
            output_path: 报告输出路径
            
        Returns:
            bool: 报告生成是否成功
        """
        try:
            # 解析扫描结果
            summary = self._parse_scan_results(scan_result)
            
            # 生成报告内容
            report_content = self._generate_report_content(summary, image_name)
            
            # 确保输出目录存在
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # 写入报告文件
            with open(output_path, "w", encoding="utf-8") as f:
                f.write(report_content)
            
            return True
            
        except Exception as e:
            print(f"生成安全报告时发生错误: {e}")
            return False

    def _parse_scan_results(self, scan_result: Dict) -> Dict:
        """解析扫描结果，提取关键信息"""
        summary = {
            "total_vulnerabilities": 0,
            "critical": 0,
            "high": 0,
            "medium": 0,
            "low": 0,
            "vulnerabilities": [],
            "os_packages": [],
            "language_packages": []
        }
        
        if not scan_result or "Results" not in scan_result:
            return summary
        
        for result in scan_result.get("Results", []):
            result_type = result.get("Type", "")
            vulnerabilities = result.get("Vulnerabilities", [])
            
            for vuln in vulnerabilities:
                severity = vuln.get("Severity", "UNKNOWN").upper()
                
                # 统计数量
                summary["total_vulnerabilities"] += 1
                if severity == "CRITICAL":
                    summary["critical"] += 1
                elif severity == "HIGH":
                    summary["high"] += 1
                elif severity == "MEDIUM":
                    summary["medium"] += 1
                elif severity == "LOW":
                    summary["low"] += 1
                
                # 收集详细信息
                vuln_info = {
                    "id": vuln.get("VulnerabilityID", ""),
                    "pkg_name": vuln.get("PkgName", ""),
                    "installed_version": vuln.get("InstalledVersion", ""),
                    "fixed_version": vuln.get("FixedVersion", ""),
                    "severity": severity,
                    "title": vuln.get("Title", ""),
                    "description": vuln.get("Description", "")[:200] + "..." if len(vuln.get("Description", "")) > 200 else vuln.get("Description", ""),
                    "type": result_type
                }
                summary["vulnerabilities"].append(vuln_info)
                
                # 分类存储
                if result_type in ["ubuntu", "alpine", "centos", "rhel", "opensuse"]:
                    summary["os_packages"].append(vuln_info)
                else:
                    summary["language_packages"].append(vuln_info)
        
        return summary

    def _generate_report_content(self, summary: Dict, image_name: str) -> str:
        """生成报告内容"""
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        
        # 安全等级评估
        security_level = self._assess_security_level(summary)
        
        report = f"""# 容器镜像安全扫描报告

## 基本信息
- **镜像名称**: {image_name}
- **扫描时间**: {timestamp}
- **安全等级**: {security_level}

## 漏洞统计
- **总漏洞数**: {summary['total_vulnerabilities']}
- **严重 (Critical)**: {summary['critical']}
- **高危 (High)**: {summary['high']}
- **中危 (Medium)**: {summary['medium']}
- **低危 (Low)**: {summary['low']}

## 安全建议
{self._generate_security_recommendations(summary)}

## 详细漏洞列表

### 严重和高危漏洞
"""
        
        # 添加严重和高危漏洞详情
        critical_high = [v for v in summary['vulnerabilities'] 
                        if v['severity'] in ['CRITICAL', 'HIGH']]
        
        if critical_high:
            for vuln in critical_high[:10]:  # 限制显示前10个
                report += f"""
#### {vuln['id']} - {vuln['severity']}
- **组件**: {vuln['pkg_name']} ({vuln['installed_version']})
- **修复版本**: {vuln.get('fixed_version', '暂无')}
- **标题**: {vuln['title']}
- **描述**: {vuln['description']}
"""
        else:
            report += "\n✅ 未发现严重或高危漏洞\n"
        
        # 添加统计表格
        report += f"""
## 按组件类型统计
- **操作系统包**: {len(summary['os_packages'])} 个漏洞
- **语言包/依赖**: {len(summary['language_packages'])} 个漏洞

## 修复建议
1. 优先修复严重和高危漏洞
2. 更新基础镜像到最新版本
3. 升级存在漏洞的组件包
4. 定期进行安全扫描

---
*报告由 Go容器化工具自动生成*
"""
        
        return report

    def _assess_security_level(self, summary: Dict) -> str:
        """评估安全等级"""
        critical = summary['critical']
        high = summary['high']
        medium = summary['medium']
        
        if critical > 0:
            return "🔴 高风险"
        elif high >= 5:
            return "🟡 中风险"
        elif high > 0 or medium >= 10:
            return "🟡 中风险"
        elif medium > 0:
            return "🟢 低风险"
        else:
            return "✅ 安全"

    def _generate_security_recommendations(self, summary: Dict) -> str:
        """生成安全建议"""
        recommendations = []
        
        if summary['critical'] > 0:
            recommendations.append("⚠️  发现严重漏洞，建议立即修复")
        
        if summary['high'] >= 3:
            recommendations.append("📋 高危漏洞较多，建议优先处理")
        
        if len(summary['os_packages']) > len(summary['language_packages']):
            recommendations.append("🔄 建议更新基础镜像版本")
        
        if summary['total_vulnerabilities'] > 50:
            recommendations.append("📊 漏洞总数较多，建议系统性修复")
        
        if not recommendations:
            recommendations.append("✅ 当前镜像安全状况良好")
        
        return "\n".join([f"- {rec}" for rec in recommendations])

    def scan_and_report(self, image_name: str, report_path: str = None) -> tuple[bool, str]:
        """
        执行完整的扫描和报告生成流程
        
        Args:
            image_name: 镜像名称
            report_path: 报告输出路径，为None时自动生成
            
        Returns:
            tuple: (成功标志, 报告路径)
        """
        if report_path is None:
            safe_name = image_name.replace(":", "_").replace("/", "_")
            report_path = f"./{safe_name}_security_report.md"
        
        print(f"🔍 正在扫描镜像安全漏洞: {image_name}")
        
        # 执行扫描
        scan_result = self.scan_image(image_name)
        if scan_result is None:
            return False, ""
        
        print("📝 正在生成安全报告...")
        
        # 生成报告
        success = self.generate_security_report(scan_result, image_name, report_path)
        if success:
            # 输出简要统计
            summary = self._parse_scan_results(scan_result)
            print(f"   发现漏洞: 严重({summary['critical']}) 高危({summary['high']}) 中危({summary['medium']}) 低危({summary['low']})")
            return True, report_path
        else:
            return False, ""


# 使用 Docker 运行 trivy 的备选方案
def scan_with_docker_trivy(image_name: str) -> Optional[Dict]:
    """使用 Docker 容器运行 trivy 扫描（当本地未安装 trivy 时的备选方案）"""
    try:
        cmd = [
            "docker", "run", "--rm",
            "-v", "/var/run/docker.sock:/var/run/docker.sock",
            "aquasec/trivy:latest",
            "image", "--format", "json",
            "--severity", "HIGH,CRITICAL,MEDIUM,LOW",
            image_name
        ]
        
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        
        if result.returncode != 0:
            print(f"Docker trivy 扫描失败: {result.stderr}")
            return None
            
        return json.loads(result.stdout) if result.stdout else {}
        
    except Exception as e:
        print(f"Docker trivy 扫描出错: {e}")
        return None